import cv2
import numpy as np
from pyapriltags import Detector

class CameraCalibrator:
    def __init__(self, rtsp_url):
        self.rtsp_url = rtsp_url
        self.detector = Detector(families="tag36h11")
    
    def calibrate(self, known_distance=1.0, real_width=0.13):
        """执行相机标定"""
        cap = cv2.VideoCapture(self.rtsp_url)
        if not cap.isOpened():
            return {"error": "无法连接摄像头"}
        
        ret, frame = cap.read()
        if not ret:
            cap.release()
            return {"error": "无法读取摄像头画面"}
        
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        results = self.detector.detect(gray)
        
        if not results:
            cap.release()
            return {"error": "未检测到标记"}
        
        r = results[0]
        corners = r.corners
        width_px = np.linalg.norm(corners[1] - corners[0])
        
        # 计算焦距
        focal_length = (width_px * known_distance) / real_width
        
        # 创建相机矩阵
        height, width = frame.shape[:2]
        camera_matrix = np.array([
            [focal_length, 0, width/2],
            [0, focal_length, height/2],
            [0, 0, 1]
        ], dtype=np.float32)
        
        # 保存参数
        dist_coeffs = np.zeros((4, 1))
        np.savez('camera_params.npz', 
                 camera_matrix=camera_matrix, 
                 dist_coeffs=dist_coeffs,
                 real_width=real_width)
        
        cap.release()
        return {
            "focal_length": float(focal_length),
            "camera_matrix": camera_matrix.tolist(),
            "status": "标定成功"
        }